from pandas import read_csv
from datetime import datetime
from climate_indices import indices
import os
import numpy as np
import math
import datetime

# load data
def parse(x):
    return datetime.strptime(x, '%Y %m')

all_indices_root = os.path.abspath('result/qinghai')  # 临时站点数据存放
spei_root = os.path.abspath('result/qinghai/example')

if not os.path.exists(spei_root):
    os.mkdir(spei_root)
multi_spei = os.path.abspath('result/qinghai/example/SPEI-12')# 复合SPEI值存放
if not os.path.exists(multi_spei):
    os.mkdir(multi_spei)
list = os.listdir(all_indices_root)

for i in range(0, len(list)):

    path = os.path.join(all_indices_root, list[i])
    if os.path.isfile(path):
        dataframe = read_csv(path, header=None, names=('year', 'month','latitude','precipitation','average_air_pressure',
                                                       'temperature','average_water_air_pressure','humidity','low_temp','high_temp'))
        precips_np = dataframe['precipitation'].values
        tem_np = dataframe['temperature'].values
        start_year = int(dataframe.iloc[0, 0])
        end_year = int(dataframe.iloc[-1, 0])
        latitude = float(dataframe.iloc[0, 2])
        distribution = indices.Distribution.gamma
        spei_gm = indices.spei(scale=12, distribution=distribution, periodicity='monthly', data_start_year=start_year,
                               calibration_year_initial=start_year, calibration_year_final=end_year,
                               precips_mm=precips_np,
                               temps_celsius=tem_np, latitude_degrees=latitude, pet_mm=None)

        distribution = indices.Distribution.pearson_type3
        spei_p3 = indices.spei(scale=12, distribution=distribution, periodicity='monthly', data_start_year=start_year,
                            calibration_year_initial=start_year, calibration_year_final=end_year, precips_mm=precips_np,
                            temps_celsius=tem_np, latitude_degrees=latitude, pet_mm=None)

        col_name = dataframe.columns.tolist()  # 将数据框的列名全部提取出来存放在列表里

        col_name.insert(2, 'spei_gamma')  # 在列索引为2的位置插入一列,列名为:spei，刚插入时不会有值，整列都是NaN
        col_name.insert(3, 'spei_pearson3')
        dataframe = dataframe.reindex(columns=col_name)  # DataFrame.reindex() 对原行/列索引重新构建索引值
        dataframe['spei_gamma']=spei_gm
        dataframe['spei_pearson3'] = spei_p3
        dataframe = dataframe[0:]
        path = os.path.join(multi_spei, 'Multi_SPEI-12_' + list[i][14:19] + '.txt')
        dataframe.to_csv(path,index =0,columns=['year', 'month', 'spei_gamma','spei_pearson3','latitude', 'precipitation','average_air_pressure',
                                                       'temperature','average_water_air_pressure','humidity','low_temp', 'high_temp'])



        # col_name.insert(2, 'spei')  # 在列索引为2的位置插入一列,列名为:spei，刚插入时不会有值，整列都是NaN
        # dataframe = dataframe.reindex(columns=col_name)  # DataFrame.reindex() 对原行/列索引重新构建索引值
        # dataframe['spei']=spei
        # dataframe = dataframe[0:]
        # path = os.path.join(multi_spei, 'Multi_SPEI-12_' + list[i][14:19] + '.txt')
        # dataframe.to_csv(path,index =0,columns=['year', 'month', 'spei', 'precipitation','average_air_pressure',
        #                                                'temperature','average_water_air_pressure','humidity','low_temp', 'high_temp'])
